"""
gRPC Agent Adapter
将现有的白方智能体接入基于gRPC的仿真系统

没做健壮性处理
锁定指定的单位

作者: 何逸凡
版本: 1.0
"""

# TODO 
# 1. 修复uav/usv死掉的状态
# 2. 检查记忆能否成功  
# 3. 为什么命令没有执行？ 无人机无法降落？
# 4. 检查无人机能不能降落

import sys
import grpc
import json
import math
import time
import re
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any

# 将项目根目录添加到Python路径，以便导入模块
sys.path.insert(0, sys.path[0]+"/../")


import os
# os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

# 仿真系统通信协议
from sh15.simulation import simserver_pb2
from sh15.simulation import simserver_pb2_grpc

# 智能体和观测/动作定义
from sh15.agents.fsm_defense_agent import create_fsm_defense_agent
from sh15.agents.base_agent import AgentObservation, AgentAction

# --- 适配器所需的数据类 ---
# 这些类模拟了原始仿真环境中的单位对象，使得智能体无需修改即可使用
@dataclass
class MockUnit:
    id: int
    x: float
    y: float
    name: str
    unit_type: str
    active: bool = True
    heading: float = 0.0  # 弧度, 0=正东方向, 逆时针为正
    speed: float = 0.0
    detected_targets: list = field(default_factory=list)
    freeze_until: float = 0.0  # 冻结结束时间
    
    def distance_to(self, other) -> float:
        """计算到另一个单位的欧几里得距离"""
        return math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
    
    def is_frozen(self, current_time: float) -> bool:
        """检查单位是否处于冻结状态"""
        return current_time < self.freeze_until

@dataclass
class MockWhiteBoat(MockUnit):
    is_locking: bool = False
    locked_times: int = 0
    locking_unit: Optional[str] = None
    unit_type: str = "white_boat"
    # 智能体依赖的额外属性
    current_target: Optional[Any] = None
    lock_start_time: float = 0.0
    max_speed: float = 10.0  # 白方无人艇最大速度
    min_turn_radius: float = 500.0  # 最小转弯半径
    
    def is_frozen(self, current_time: float) -> bool:
        """重写：从接口的is_frozen字段判断"""
        # 注意：freeze_until的值在对象创建时根据is_frozen标志设置
        return current_time < self.freeze_until

@dataclass
class MockWhiteDrone(MockUnit):
    energy: float = 100.0  # gRPC接口中是剩余飞行秒数'battery'，这里转换为百分比
    state: str = "recharging"  # flying, returning, recharging, ready
    docked_on_boat_id: Optional[int] = -1
    parent_boat_id: Optional[int] = -1 # 需要根据name推断
    unit_type: str = "white_drone"
    max_speed: float = 150.0  # 白方无人机最大速度
    min_speed: float = 20.0   # 白方无人机最小速度
    min_turn_radius: float = 1000.0  # 最小转弯半径
    # 以“百分比/分钟”为单位的能量变化速率，假设满电可飞行120分钟、充电也按120分钟充满
    energy_consumption_rate: float = 100.0 / 120.0
    energy_recharge_rate: float = 100.0 / 120.0
    
    # 原始对象中包含的额外属性，以防万一
    is_charging: bool = False
    is_at_usv: bool = True
    home_name: Optional[str] = None


@dataclass
class MockBlackBoat(MockUnit):
    locked_times: int = 0
    unit_type: str = "black_boat"
    crossed: bool = False  # 是否已突破防线
    max_speed: float = 10.0  # 黑方无人艇最大速度
    # 模拟原始对象，添加velocity属性
    velocity: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
    is_frozen: bool = False


# --- 坐标系与单位转换工具 ---
def grpc_course_to_agent_heading(course_deg: float) -> float:
    """
    将gRPC接口的航向 (度, 0=正北, 顺时针) 转换为智能体使用的航向 (弧度, 0=正东, 逆时针).
    """
    # 1. 将gRPC航向转换为数学常规角度 (0=正东, 逆时针)
    # gRPC的0度(北)是数学的90度。gRPC顺时针增加，数学逆时针增加。
    # 所以, math_deg = 90 - grpc_deg
    math_deg = (90 - course_deg) % 360
    # 2. 将角度转换为弧度
    return math.radians(math_deg)

def agent_heading_to_grpc_course(heading_rad: float) -> float:
    """
    将智能体使用的航向 (弧度, 0=正东, 逆时针) 转换为gRPC接口的航向 (度, 0=正北, 顺时针).
    """
    # 1. 将弧度转换为角度
    math_deg = math.degrees(heading_rad)
    # 2. 将数学角度转换为gRPC航向
    # grpc_deg = 90 - math_deg
    grpc_deg = (90 - math_deg) % 360
    return grpc_deg
    
def parse_unit_id(unit_name: str) -> int:
    """从单位名称中解析出数字ID"""
    if not unit_name:
        return -1
    match = re.search(r'\d+', unit_name)
    return int(match.group()) if match else -1


class GrpcAgentAdapter:
    def __init__(self, host='127.0.0.1', port='6000', user_name='admin'):
        """初始化适配器"""
        # gRPC客户端设置
        self.channel = grpc.insecure_channel(f"{host}:{port}")
        self.client = simserver_pb2_grpc.GreeterStub(channel=self.channel)
        self.user_name = user_name
        print("gRPC客户端已连接。")

        # 智能体初始化
        self.agent = create_fsm_defense_agent()
        self.agent.reset()
        print("白方FSM智能体已初始化。")
        
        # 存储单位名称到ID的映射
        self.white_boat_id_map: Dict[int, str] = {}
        self.white_drone_id_map: Dict[int, str] = {}
        self.black_boat_id_map: Dict[int, str] = {}
        # 作为外部“记忆”，存储所有已知黑方单位
        self.known_black_boats: Dict[int, MockBlackBoat] = {}
        # 上次状态更新时间戳
        self.last_update_time: float = 0.0
        # 缓存无人机母船ID映射（0基）
        self.drone_parent_boat_map: Dict[int, int] = {}
        
        # 缓存当前帧的mock对象，供_execute_actions使用
        self.current_mock_boats: List[MockWhiteBoat] = []
        self.current_mock_drones: List[MockWhiteDrone] = []

    def _send_grpc_request(self, func_name: str, **kwargs) -> Optional[Dict]:
        """发送gRPC请求并返回JSON解析后的结果"""
        try:
            msg = json.dumps({
                "source": "test",
                "ip":"test",
                "user_name": self.user_name,
                "engine_name":"",
                "flag": "simulation",
                "func_name": func_name,
                "kwargs": kwargs or {}
            })
            request = self.client.control(simserver_pb2.MsgStr(msg=msg))  # type: ignore
            if request.data:
                return json.loads(request.data)
            # 对于不返回data的指令，可以检查msg
            if "success" in request.msg or "True" in request.msg:
                return {"status": "success"}
            return None
        except grpc.RpcError as e:
            print(f"gRPC请求失败: {func_name}, 错误: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"解析gRPC响应失败: {func_name}, 响应: {request.data}, 错误: {e}")
            return None


    def run(self):
        """运行主循环"""
        print("启动智能体决策循环...")
        while True:
            # 1. 获取仿真状态
            state = self._send_grpc_request("get_state")
            if not state:
                # 如果获取状态失败，可以稍等一下再重试
                time.sleep(0.1)
                print("获取状态失败，等待0.1秒后重试...")
                continue
            else:
                # print(state)
                pass

            # 2. 转换状态为智能体所需的格式
            current_time = state.get("time", 0.0)
            dt = current_time - self.last_update_time
            self.last_update_time = current_time

            mock_white_boats = self._create_mock_white_boats(state.get("white_usv_states", []), current_time)
            mock_white_drones = self._create_mock_white_drones(state.get("white_uav_states", []))
            
            # 更新并预测黑方单位位置
            all_known_black_boats = self._update_and_predict_black_boats(state, dt)

            # 缓存mock对象，供_execute_actions使用
            self.current_mock_boats = mock_white_boats
            self.current_mock_drones = mock_white_drones

            # 3. 创建智能体观测对象
            # 类型忽略：Mock类实现了所有必要的接口
            observation = AgentObservation(
                white_boats=mock_white_boats,  # type: ignore
                all_drones=mock_white_drones,  # type: ignore
                black_boats=all_known_black_boats,  # type: ignore
                current_time=current_time
            )
            # print(mock_white_drones)
            # 4. 获取智能体动作
            actions = self.agent.get_actions(observation)
            #print(actions)
            # 5. 执行动作
            # print(actions)
            self._execute_actions(actions)
            
            # 根据用户要求，不添加延时，以最高速度运行

    def _create_mock_white_boats(self, white_usv_states: List[Dict], current_time: float = 0.0) -> List[MockWhiteBoat]:
        boats = []
        for s in white_usv_states:
            uid = parse_unit_id(s["name"])
            if uid == -1: continue

            # 将外部1基ID转换为内部0基索引
            id_zero_based = max(0, uid - 1)
            self.white_boat_id_map[id_zero_based] = s["name"]
            
            # 增加is_alive检查
            if not s["is_alive"]:
                boats.append(MockWhiteBoat(
                    id=id_zero_based, 
                    name=s["name"], 
                    active=False, 
                    x=s.get("position", [0, 0])[0], 
                    y=s.get("position", [0, 0])[1]
                ))
                continue
            
            # 处理freeze_until：如果is_frozen为True，设置一个大的未来时间
            freeze_until = current_time + 300.0 if s["is_frozen"] else 0.0

            boat = MockWhiteBoat(
                id=id_zero_based,
                x=s["position"][0],
                y=s["position"][1],
                name=s["name"],
                active=s["is_alive"],
                speed=s["speed"],
                heading=grpc_course_to_agent_heading(s["course"]),
                is_locking=s["is_locking"],
                locked_times=s["locked_times"],
                locking_unit=s.get("locking_unit"),
                freeze_until=freeze_until
            )
            boats.append(boat)
        return boats
        
    def _create_mock_white_drones(self, white_uav_states: List[Dict]) -> List[MockWhiteDrone]:
        drones = []
        for s in white_uav_states:
            uid = parse_unit_id(s["name"])
            if uid == -1: continue

            # 将外部1基ID转换为内部0基索引
            id_zero_based = max(0, uid - 1)
            self.white_drone_id_map[id_zero_based] = s["name"]
            
            # 增加is_alive检查
            if not s["is_alive"]:
                drones.append(MockWhiteDrone(
                    id=id_zero_based, 
                    name=s["name"], 
                    active=False, 
                    x=s.get("position", [0,0])[0], 
                    y=s.get("position", [0,0])[1]
                ))
                continue
            
            # 判断无人机状态
            docked_id = -1
            if s["is_at_usv"]:
                # 停靠在母舰上
                home_uid = parse_unit_id(s["home_name"]) if s.get("home_name") else -1
                docked_id = max(0, home_uid - 1) if home_uid != -1 else -1
                # 区分充电和就绪状态
                state = "recharging" if s["is_charging"] else "ready"
                # === 新增：缓存母船ID ===
                if docked_id != -1:
                    self.drone_parent_boat_map[id_zero_based] = docked_id
            else:
                # 在空中飞行
                state = "flying"
            
            # === 修改：优先从缓存中读取母船ID ===
            parent_id = self.drone_parent_boat_map.get(id_zero_based, -1)
            # Fallback for safety, should not happen in normal flow
            if parent_id == -1:
                 parent_id = docked_id if docked_id >= 0 else (max(0, parse_unit_id(s.get("home_name", "")) - 1) if s.get("home_name") else -1)

            drone = MockWhiteDrone(
                id=id_zero_based,
                x=s["position"][0],
                y=s["position"][1],
                name=s["name"],
                active=s["is_alive"],
                speed=s["speed"],
                heading=grpc_course_to_agent_heading(s["course"]),
                energy=(s["battery"] / (2 * 3600)) * 100, # 假设满电2小时，转换为百分比
                state=state,
                docked_on_boat_id=docked_id,
                parent_boat_id=parent_id,
                is_charging=s["is_charging"],
                is_at_usv=s["is_at_usv"],
                home_name=s["home_name"]
            )
            drones.append(drone)
        return drones

    def _update_and_predict_black_boats(self, state: Dict, dt: float) -> List[MockBlackBoat]:
        """
        更新已知黑方单位状态，并为未观测到的单位预测位置
        """
        # 1. 解析当前帧的观测，获取当前可见的黑方单位的最新信息
        current_frame_detected_boats = self._create_mock_black_boats(
            state.get("white_observation", []),
            state.get("black_usv_states", [])
        )
        seen_ids = set(current_frame_detected_boats.keys())

        # 2. 预测未观测到的船只位置
        if dt > 0: # 只有在时间前进时才预测
            for boat_id, boat in self.known_black_boats.items():
                if boat_id not in seen_ids and boat.active:
                    # 简单线性预测
                    boat.x += boat.velocity[0] * dt
                    boat.y += boat.velocity[1] * dt
        
        # 3. 更新或添加新观测到的船只信息到知识库
        for boat_id, boat in current_frame_detected_boats.items():
            self.known_black_boats[boat_id] = boat
        
        # 4. 返回所有已知（包括预测的）黑方单位的列表
        return list(self.known_black_boats.values())

    def _create_mock_black_boats(self, white_observation: List[Dict], black_usv_states: List[Dict]) -> Dict[int, MockBlackBoat]:
        """
        仅从当前帧的接口数据创建MockBlackBoat对象字典，不进行状态管理。
        """
        # 1. 从 black_usv_states 创建状态查找字典 {id: state_dict}
        black_states = {}
        for s in black_usv_states:
            uid = parse_unit_id(s["name"])
            if uid != -1:
                id_zero_based = max(0, uid - 1)
                black_states[id_zero_based] = {
                    "active": s["is_alive"],
                    "locked_times": s["locked_times"],
                    "is_frozen": s["is_frozen"]
                }
                # 同时更新ID映射
                self.black_boat_id_map[id_zero_based] = s["name"]

        # 2. 遍历白方视野，创建信息完整的Mock对象
        detected_boats_map = {}
        for s in white_observation:
            uid = parse_unit_id(s["name"])
            if uid == -1: continue
            
            id_zero_based = max(0, uid - 1)
            
            # 获取单位状态，如果没在全局状态列表中找到，则使用默认值
            state = black_states.get(id_zero_based, {
                "active": True, "locked_times": 0, "is_frozen": False
            })

            # 计算运动学信息
            vx, vy = s["velocity"][0], s["velocity"][1]
            speed = math.sqrt(vx**2 + vy**2)
            heading = math.atan2(vy, vx) if speed > 0.01 else 0.0

            # 创建包含完整信息的Mock对象
            boat = MockBlackBoat(
                id=id_zero_based,
                x=s["position"][0],
                y=s["position"][1],
                name=s["name"],
                active=state["active"],
                speed=speed,
                heading=heading,
                velocity=[vx, vy, 0], # 填充为三维向量
                locked_times=state["locked_times"],
                is_frozen=state["is_frozen"]
            )
            detected_boats_map[id_zero_based] = boat
            
        return detected_boats_map
        
    def _execute_actions(self, actions: AgentAction):
        """将AgentAction转换为gRPC调用"""
        raw_boat_actions = actions.boat_actions
        raw_drone_actions = actions.drone_actions

        # 将动作数组转换为以ID为键的字典，提高查找效率和安全性
        boat_actions = {i: raw_boat_actions[i] for i in range(len(raw_boat_actions))}
        drone_actions = {i: raw_drone_actions[i] for i in range(len(raw_drone_actions))}
        
        # 执行无人艇动作
        # 直接遍历当前帧的mock对象，使用其ID作为actions数组索引
        for boat in self.current_mock_boats:
            if not boat.active:
                continue
            
            # 使用字典get方法安全获取动作
            action = boat_actions.get(boat.id)
            if action is None:
                print(f"[警告] 无法找到无人艇ID {boat.id} 的对应动作")
                continue
            
            # 获取该无人艇的动作
            speed, heading, lock_id = action
            unit_name = boat.name  # 直接使用mock对象的name字段
            
            # 1. 发送移动指令
            target_course = agent_heading_to_grpc_course(heading)
            self._send_grpc_request(
                "send_command",
                cmd={
                    "unit_name": unit_name,
                    "target_speed": float(speed),
                    "target_course": float(target_course),
                }
            )
            
            # 2. 发送锁定指令
            if lock_id >= 0:
                target_name = self.black_boat_id_map.get(int(lock_id))
                if target_name:
                    self._send_grpc_request("cmd_lock", unit1_name=unit_name, unit2_name=target_name)
                # else:
                #     print(f"[警告] 无法找到黑方目标 ID={int(lock_id)} 的名称")
            # 注意：接口没有提供"停止锁定"的指令，通常锁定会在超时或目标超出范围后自动解除。

        # 执行无人机动作
        if drone_actions:
            for drone in self.current_mock_drones:
                if not drone.active:
                    continue
                
                # 使用字典get方法安全获取动作
                action = drone_actions.get(drone.id)
                if action is None:
                    print(f"[警告] 无法找到无人机ID {drone.id} 的对应动作")
                    continue
                
                # 获取该无人机的动作
                speed, heading, drone_action_val = action
                unit_name = drone.name  # 直接使用mock对象的name字段
                
                # 1. 移动指令（对所有状态都发送，仿真系统会处理是否能移动）
                target_course = agent_heading_to_grpc_course(heading)
                # print({
                #         "unit_name": unit_name,
                #         "target_speed": float(speed),
                #         "target_course": float(target_course),
                #     })
                self._send_grpc_request(
                    "send_command",
                    cmd={
                        "unit_name": unit_name,
                        "target_speed": float(speed),
                        "target_course": float(target_course),
                    }
                )
                # 2. 起飞/降落指令
                # 逻辑复刻自 ActionExecutor
                # drone_action < 0.5 降落, >= 0.5 起飞/飞行
                if drone_action_val < 0.5: # 降落
                    # 从无人机对象的home_name字段获取母舰名称
                    home_name = drone.home_name
                    if home_name and drone.state in ["flying", "returning"]:
                        # 只有在飞行状态才发送降落指令
                        self._send_grpc_request("cmd_uav_land", uav_name=unit_name, usv_name=home_name)
                else: # 起飞/飞行
                    # 如果无人机在母舰上（就绪或正在充电），尝试起飞
                    if drone.state in ["recharging", "ready"] and drone.docked_on_boat_id is not None and drone.docked_on_boat_id >= 0:
                        home_name = drone.home_name
                        if home_name:
                            # 起飞时也需要速度和航向
                            takeoff_course = agent_heading_to_grpc_course(heading)
                            self._send_grpc_request("cmd_uav_takeoff", unit_name=unit_name, home_name=home_name, target_speed=float(speed), target_course=float(takeoff_course))
                    # 如果已经在飞行，移动指令已经在上面发送了


if __name__ == "__main__":
    adapter = GrpcAgentAdapter()
    try:
        adapter.run()
    except KeyboardInterrupt:
        print("\n用户中断程序。")
    except Exception as e:
        print(f"\n程序出现严重错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        print("程序退出。")
